import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import pinv

# ========================== 0. 从 TXT 文件中读取数据，自动确定 dt 和 tau ==========================
txt_path = "uav1的数据.txt"  # 文件名（需与脚本在同一目录，或使用完整路径）

# 手动读取首行并替换中文逗号为英文逗号
with open(txt_path, "r", encoding="utf-8") as f:
    lines = f.readlines()

# 替换首行分隔符
lines[0] = lines[0].replace("，", ",")  # 替换全角逗号为英文逗号

# 将处理后的数据写入一个缓存中以便 pd.read_csv 使用
from io import StringIO
cleaned_data = StringIO("".join(lines))

# 使用 pandas 读取修正后的数据
df = pd.read_csv(cleaned_data)

# 去除列名空格
df.columns = df.columns.str.strip()
print("列名列表：", df.columns.tolist())

# 提取时间轴
time_array = df["Time"].to_numpy()
dt_array = np.diff(time_array)

if len(dt_array) == 0:
    raise ValueError("时间列不足两个数据点，无法计算采样间隔 dt")

dt = np.mean(dt_array)
T = time_array[-1] - time_array[0]
tau = len(time_array)
print(f"[INFO] 采样间隔 dt = {dt:.6f} s, 总时长 T = {T:.3f} s, 样本数 τ = {tau}")

# ========================== 1. 从 txt 中读取球坐标测量数据 ==========================
def generate_leader_measurements():
    try:
        r = df["斜距"].to_numpy()
        elev = df["高低角"].to_numpy()  # 高低角（弧度）
        azim = df["方位角"].to_numpy()  # 方位角（弧度）
        azim_dot = df["方位角变化率"].to_numpy()  # 方位角变化率
    except KeyError as e:
        raise KeyError("列名读取失败，请确保数据列顺序为：Time, 序号, 斜距, 高低角, 方位角, 方位角变化率") from e

    # 计算斜距变化率和高低角变化率（数值微分）
    r_dot = np.gradient(r, dt)  # 斜距变化率
    elev_dot = np.gradient(elev, dt)  # 高低角变化率
    # 构造测量数据矩阵（与之前格式一致）
    meas = np.vstack([time_array, np.zeros_like(time_array), r, elev, azim, r_dot, elev_dot, azim_dot]).T
    return meas

# ========================== 2. 球坐标转笛卡尔 ==========================
def spherical_to_cartesian(meas):
    _, _, r, elev, azim, r_dot, elev_dot, azim_dot = meas.T
    # 位置计算（x, y, z）
    x = r * np.cos(elev) * np.cos(azim)
    y = r * np.cos(elev) * np.sin(azim)
    z = r * np.sin(elev)

    # 速度计算（x_dot, y_dot, z_dot）
    x_dot = (r_dot * np.cos(elev) * np.cos(azim) 
             - r * elev_dot * np.sin(elev) * np.cos(azim) 
             - r * azim_dot * np.cos(elev) * np.sin(azim))
    y_dot = (r_dot * np.cos(elev) * np.sin(azim) 
             - r * elev_dot * np.sin(elev) * np.sin(azim) 
             + r * azim_dot * np.cos(elev) * np.cos(azim))
    z_dot = r_dot * np.sin(elev) + r * elev_dot * np.cos(elev)

    # 返回状态：[x, x_dot, y, y_dot, z, z_dot]
    return np.vstack([x, x_dot, y, y_dot, z, z_dot]).T

# ========================== 3. 生成持续激励输入 ==========================
def generate_input():
    # 生成随机输入序列（满足持续激励条件）
    return np.random.randn(tau - 1, 3) * 2  # 3维输入：ax, ay, az

# ========================== 4. 数据驱动估计 S1 ==========================
def estimate_S1(leader_state, U):
    gamma = leader_state.T  # γ ∈ R^{6×τ}（6维状态，τ个时刻）
    gamma_L = gamma[:, :-1]  # γ(k)，k=0到τ-2
    gamma_R = gamma[:, 1:]   # γ(k+1)，k=1到τ-1
    U_L = U.T  # 控制输入序列，shape=(3, τ-1)

    # 增广数据矩阵 M = [γ_L; U_L] ∈ R^{9×(τ-1)}
    M = np.vstack([gamma_L, U_L])
    rank_M = np.linalg.matrix_rank(M)
    print(f"[INFO] 增广矩阵秩 Rank(M) = {rank_M}, 理论需求 = 9")

    # 估计 S1 和 B_i1
    S_all = gamma_R @ pinv(M)  # 利用伪逆求解
    S1 = S_all[:, :6]   # 提取 S1（前6列）
    B_i1 = S_all[:, 6:] # 提取 B_i1（后3列）
    
    return S1, B_i1

# ========================== 5. 状态递推仿真（基于估计的 S1） ==========================
def simulate_leader_from_S1(S1, v0, steps):
    n = S1.shape[0]
    v = np.zeros((steps + 1, n))  # 存储状态序列
    v[0] = v0  # 初始状态
    for k in range(steps):
        v[k + 1] = S1 @ v[k]  # 递推公式：v(k+1) = S1 * v(k)
    return v

# ========================== 6. 可视化函数 ==========================
def plot_3d_trajectory(positions, title, subplot_pos):
    ax = plt.subplot(subplot_pos, projection='3d')
    ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], label="Trajectory")
    ax.set_xlabel("X (m)")
    ax.set_ylabel("Y (m)")
    ax.set_zlabel("Z (m)")
    ax.set_title(title)
    ax.legend()

def visualize(S1, leader_state, v_pred, time_array):
    # 1. S1特征值
    eigs = np.linalg.eigvals(S1)
    plt.figure(figsize=(18, 5))
    plt.subplot(131)
    plt.scatter(np.real(eigs), np.imag(eigs), marker='x', color='r')
    plt.title("S1 Eigenvalue Distribution")
    plt.xlabel("Real Part")
    plt.ylabel("Imaginary Part")
    plt.grid(True)
    plt.axhline(0, color='k', ls='--', alpha=0.3)
    plt.axvline(0, color='k', ls='--', alpha=0.3)

    # 2. 真实轨迹与估计轨迹（数据时段）
    plot_3d_trajectory(leader_state[:, [0, 2, 4]], "True Trajectory", 132)
    plot_3d_trajectory(v_pred[:tau, [0, 2, 4]], "Estimated Trajectory", 133)
    plt.tight_layout()
    plt.show()

    # 3. 未来轨迹预测
    plt.figure()
    plot_3d_trajectory(v_pred[tau:, [0, 2, 4]], "Future Trajectory", 111)
    plt.tight_layout()
    plt.show()

    # 4. 误差分析
    error_vec = v_pred[:tau, [0, 2, 4]] - leader_state[:, [0, 2, 4]]  # 位置误差
    error_norm = np.linalg.norm(error_vec, axis=1)  # 误差模长
    plt.figure()
    plt.plot(time_array, error_norm, label="Position Error")
    plt.xlabel("t (s)")
    plt.ylabel("Position Error (m)")
    plt.title("Trajectory Prediction Error During Data Period")
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.show()

# ========================== 7. 主程序入口 ==========================
def main():
    # 读取测量数据
    meas = generate_leader_measurements()
    # 转换为笛卡尔坐标下的状态
    leader_state = spherical_to_cartesian(meas)
    # 生成控制输入
    U = generate_input()
    # 估计 S1 和 B_i1
    S1, B_i1 = estimate_S1(leader_state, U)

    # 输出估计结果
    print("\n=== 领导者系统矩阵 S1 估计结果 ===")
    print(np.round(S1, 4))
    print("\n=== 辅助系统输入矩阵 B_i1 估计结果 ===")
    print(np.round(B_i1, 4))

    # 基于 S1 仿真轨迹（包括数据时段和未来预测）
    v0 = leader_state[0, :6]  # 初始状态
    steps_pred = tau * 2  # 预测步数（数据时段的2倍）
    v_pred = simulate_leader_from_S1(S1, v0, steps_pred)

    # 可视化
    visualize(S1, leader_state, v_pred, time_array)

    # 误差统计
    error_vec = v_pred[:tau, [0, 2, 4]] - leader_state[:, [0, 2, 4]]
    error_norm = np.linalg.norm(error_vec, axis=1)
    print("\n=== 数据时段内的误差统计 ===")
    print(f"平均误差: {np.mean(error_norm):.4f} m, 最大误差: {np.max(error_norm):.4f} m")

if __name__ == "__main__":
    main()